use std::rc::Rc;
use num_traits::{Num, NumCast};
use rand::distributions::Uniform;
use rand::{Rng, thread_rng};
use crate::Tensor;

pub fn constant<T>(shape: &[usize], value: T) -> Tensor<T>
    where
        T: Copy + Num + NumCast + 'static,
{
    let data = (0..shape.iter().product())
        .map(|_| value)
        .collect::<Box<[T]>>();
    (Rc::new(data), shape.to_vec().into_boxed_slice()).into()
}

pub fn ones<T>(shape: &[usize]) -> Tensor<T>
    where
        T: Copy + Num + NumCast + 'static,
{
    constant(shape, T::one())
}

pub fn uniform<T>(shape: &[usize], a: f64, b: f64) -> Tensor<T>
    where
        T: Copy + Num + NumCast + 'static,
{
    let data = thread_rng()
        .sample_iter(Uniform::new(a, b))
        .take(shape.iter().product())
        .map(|x| T::from(x).unwrap())
        .collect::<Box<[T]>>();
    (Rc::new(data), shape.to_vec().into_boxed_slice()).into()
}

pub fn uniform_<T>(tensor: Tensor<T>, a: f64, b: f64)
    where
        T: Copy + Num + NumCast + 'static,
{
    let mut cell = (*tensor.ptr).borrow_mut();
    cell.data = thread_rng()
        .sample_iter(Uniform::new(a, b))
        .take(cell.shape.iter().product())
        .map(|x| T::from(x).unwrap())
        .collect::<Box<[T]>>()
        .into();
}

pub fn zeros<T>(shape: &[usize]) -> Tensor<T>
    where
        T: Copy + Num + NumCast + 'static,
{
    constant(shape, T::zero())
}
